from torch.utils.data import Dataset, DataLoader
import torch
from numpy import transpose
import mxnet as mx
import os
from tqdm import tqdm


class Ms1mDataset(Dataset):
    """
       MS1M cleaned via mxnet record file.
    """

    def __init__(self, path):
        path = os.path.join(path, 'train.')
        self.imgrec = mx.recordio.MXIndexedRecordIO(path + 'idx', path + 'rec', 'r')
        header, _ = mx.recordio.unpack(self.imgrec.read_idx(0))
        self.num_imgs = int(header.label[0]) - 1
        self.num_people = int(header.label[1] - header.label[0])

    def __len__(self):
        return self.num_imgs

    def __getitem__(self, idx):
        header, img = mx.recordio.unpack_img(self.imgrec.read_idx(idx + 1))
        assert img.shape == (112, 112, 3), "img is not (112, 112, 3)"
        img_ = transpose(img.copy(), axes=(2, 0, 1))
        return torch.tensor(img_, dtype=torch.float32), torch.LongTensor([int(header.label)])


if __name__ == '__main__':
    per_batch_size = 10

    train_dataset = Ms1mDataset('/home/data1/worm/datas/faces_emore')
    train_loader = DataLoader(dataset=train_dataset, batch_size=per_batch_size, num_workers=10)
    for output in tqdm(train_loader):
        pass
